import os

import torch
from natsort import natsorted

import onnxruntime
import torchvision.transforms as transforms
from PIL import Image
from dataset import TestDataset

device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")

def predict(ort_session, image_folder):
    transform = transforms.Compose([
        transforms.Resize((112, 112)),
        transforms.ToTensor(),
        transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225])
    ])

    dataset = TestDataset(image_folder, transform=transform)
    dataloader = torch.utils.data.DataLoader(dataset, batch_size=1, shuffle=False, num_workers=0)

    classes = dataset.classes

    with torch.no_grad():
        for i, (inputs, _) in enumerate(dataloader):
            ort_inputs = {ort_session.get_inputs()[0].name: inputs.numpy()}
            outputs = ort_session.run(None, ort_inputs)
            pred = torch.from_numpy(outputs[0]).squeeze()
            id = pred.argmax()
            predicted_label = classes[id]
            print("Image {}: {}".format(i, predicted_label))

""""
测试模型效果，也可作为使用的基础
"""

if __name__ == '__main__':

    """"
    测试的参数与训练保持一致
    """
    onnx_path = 'pth/best_model.onnx'
    ort_session = onnxruntime.InferenceSession(onnx_path)
    """"
    folder_path: 测试文件夹，测试文件夹下所有图片，输出类型
    """
    folder_path = "dataset/test/3"
    files = []
    lists = os.listdir(folder_path)

    predict(ort_session, folder_path)




